#!/usr/bin/python
# -*- coding: utf-8 -*-

import numpy as np
import pandas as pd
import time
import os

SPEC_NUM = 25
ROW_NUM = SPEC_NUM
COL_NUM = SPEC_NUM
N_STATES = ROW_NUM * COL_NUM  # 1维世界的宽度
ACTIONS = ['left', 'right', 'up', 'down']     # 探索者的可用动作
ACTIONS_OFF = {
    'left': -1,
    'right': 1,
    'up': -ROW_NUM,
    'down': +ROW_NUM
}
EPSILON = 0.9      # 贪婪度 greedy
ALPHA = 0.1        # 学习率
GAMMA = 0.9        # 奖励递减值
MAX_EPISODES = 13  # 最大回合数
FRESH_TIME = 0     # 移动间隔时间

def build_q_table(n_states, actions):
    table = pd.DataFrame(
        np.zeros((n_states, len(actions))),     # q_table 全 0 初始
        columns=actions,    # columns 对应的是行为名称
    )
    return table

def choose_action(env, state, q_table):
    state_actions = q_table.iloc[state, :]  # 选出这个 state 的所有 action 值
    if (np.random.uniform() > EPSILON) or (state_actions.all() == 0):  # 非贪婪 or 或者这个 state 还没有探索过 
        acts = []
        for act in ACTIONS:
            i = state + ACTIONS_OFF[act]
            if (i >= 0 and i <= N_STATES and (env >> i & 1) == 0):
                acts.append(act)
        action_name = np.random.choice(acts)
    else:
        action_name = q_table.columns[state_actions.argmax()]    # 贪婪模式
    return action_name

def get_env_feedback(S, A, end):
    # This is how agent will interact with the environment
    S_ = S + ACTIONS_OFF[A]
    if S_ == end:
        S_ = 'terminal'
        R = 1
    else:
        R = 0
    return S_, R

def update_env(env, end, S, episode, step_counter, q_table):
    # This is how environment be updated
    # env_list = ['-']*(N_STATES-1) + ['T']   # '---------T' our environment
    if S == 'terminal':
        interaction = 'Episode %s: total_steps = %s' % (episode+1, step_counter)
        print('\r{}'.format(interaction), end='')
        time.sleep(2)
        print('\r                                ', end='')
    else:
        # env_list[S] = 'o'
        # interaction = ''.join(env_list)
        # print('\r{}'.format(interaction), end='')
        os.system("clear")
        print("\033[%d;%df" % (0, 0), end="")
        print(f'Episode {episode+1}: total_steps = {step_counter}')
        for row in range(ROW_NUM):
            for col in range(COL_NUM):
                if S == row * ROW_NUM + col:
                    print(f"? ", end="")
                elif end == row * ROW_NUM + col:
                    print(f"$ ", end="")
                else:
                    print(f"{'#' if env >> (row * ROW_NUM + col) & 1 else ' '} ", end="")
            print()
        print(q_table[(q_table != 0).any(axis=1)])
        time.sleep(FRESH_TIME)


def init_environment(fname):
    env = 0
    sta = 0
    end = 0
    with open(fname, "r") as f:
        lines = f.readlines()
    for row, line in enumerate(lines):
        for col in range(len(line) - 1):
            if line[col] == "#":
                env += 1 << (row * ROW_NUM + col)
            if line[col] == "?":
                sta = row * ROW_NUM + col
            if line[col] == "O":
                end = row * ROW_NUM + col
    print(f"sta = {sta // ROW_NUM},{sta % ROW_NUM}")
    print(f"end = {end // ROW_NUM},{end % ROW_NUM}")
    return env, sta, end

def rl():
    env, sta, end = init_environment("maze_001.txt")
    q_table = build_q_table(N_STATES, ACTIONS)  # 初始 q table
    for episode in range(MAX_EPISODES):     # 回合
        step_counter = 0
        S = sta  # 回合初始位置
        is_terminated = False   # 是否回合结束
        update_env(env, end, S, episode, step_counter, q_table)    # 环境更新
        while not is_terminated:

            A = choose_action(env, S, q_table)   # 选行为
            S_, R = get_env_feedback(S, A, end)  # 实施行为并得到环境的反馈
            q_predict = q_table.loc[S, A]    # 估算的(状态-行为)值
            if S_ != 'terminal':
                q_target = R + GAMMA * q_table.iloc[S_, :].max()   #  实际的(状态-行为)值 (回合没结束)
            else:
                q_target = R     #  实际的(状态-行为)值 (回合结束)
                is_terminated = True    # terminate this episode

            q_table.loc[S, A] += ALPHA * (q_target - q_predict)  #  q_table 更新
            S = S_  # 探索者移动到下一个 state

            update_env(env, end, S, episode, step_counter+1, q_table)  # 环境更新

            step_counter += 1
    return q_table



if __name__ == "__main__":
    q_table = rl()
    print("Q-learning")
    print(q_table)
